#include "MyCLAHE.h"
#include <opencv2/highgui/highgui.hpp>

MyCLAHE::MyCLAHE(int numTilesRow,int numTilesCol,float fCliplimit)
	:m_padRowPre(0),m_padRowPost(0),
	 m_padColPre(0),m_padColPost(0)
{
	m_numTilesRow = numTilesRow;
	m_numTilesCol = numTilesCol;
	m_fCilplimit = fCliplimit;
}

void MyCLAHE::processUsingCLAHE(cv::Mat& srcImg,cv::Mat& enhanceImg)
{
	//Extend the srcImg
	cv::Mat extendedImg;
	this->extendImage(srcImg,extendedImg);
	//cv::imshow("extendedImg",extendedImg);
	cv::Mat extendedEnhanceImg = cv::Mat::zeros(extendedImg.size(),extendedImg.type());
	//Compute the mapping histogram for each tile
	std::vector<std::vector<unsigned long int>> tileMappings;
	this->computeTileMappings(extendedImg,tileMappings);
	//Interpolate the mapping histograms to get the CLAHE image
	unsigned int subRows,subCols;
	unsigned int mapIndexTop,mapIndexBottom;
	unsigned int mapIndexLeft,mapIndexRight;
	unsigned int rowRangeTop = 0;	unsigned int rowRangeBottom = 0;
	unsigned int colRangeLeft = 0;  unsigned int colRangeRight = 0;
	unsigned int extendedRows = extendedImg.rows;
	unsigned int extendedCols = extendedImg.cols;
	for (unsigned int i=0;i<=m_numTilesRow;++i)
	{
		if (i==0) //special case: top row
		{
			subRows = m_rowsTile >> 1;
			mapIndexTop = 0;
			mapIndexBottom = 0;
		} 
		else
		{
			if (i==m_numTilesRow) //special case: bottom row
			{
				subRows = m_rowsTile >> 1;
				mapIndexTop = m_numTilesRow - 1;
				mapIndexBottom = mapIndexTop;
			} 
			else
			{
				subRows = m_rowsTile;
				mapIndexTop = i - 1;
				mapIndexBottom = i;
			}
		}
		rowRangeTop = rowRangeBottom; //Locate the sub-region in row direction
		rowRangeBottom = rowRangeTop + subRows;
		if(rowRangeBottom>=extendedRows) rowRangeBottom = extendedRows-1;
		colRangeLeft = 0;
		colRangeRight = 0;
		for (unsigned int j=0;j<=m_numTilesCol;++j)
		{
			if (j==0) //special case: left col
			{
				subCols = m_colsTile >> 1;
				mapIndexLeft = 0;
				mapIndexRight = 0;
			} 
			else
			{
				if (j==m_numTilesCol)
				{
					subCols = m_colsTile >> 1;
					mapIndexLeft = m_numTilesCol-1;
					mapIndexRight = mapIndexLeft;
				} 
				else
				{
					subCols = m_colsTile;
					mapIndexLeft = j-1;
					mapIndexRight = j;
				}
			}
			colRangeLeft = colRangeRight; //Locate the sub-region in col direction
			colRangeRight = colRangeLeft + subCols;
			if(colRangeRight>=extendedCols) colRangeRight = extendedCols-1;
			//mapHist_LT mapIndexTop*m_numTilesRow+mapIndexLeft
			//mapHist_LB mapIndexBottom*m_numTilesRow+mapIndexLeft
			//mapHist_RB mapIndexBottom*m_numTilesRow+mapIndexRight
			//mapHist_RT mapIndexTop*m_numTilesRow+mapIndexLeft
			cv::Range rowRange(rowRangeTop,rowRangeBottom);
			cv::Range colRange(colRangeLeft,colRangeRight);
			cv::Mat subSrcRegion(extendedImg,rowRange,colRange);
			cv::Mat subEnhanceRegion(extendedEnhanceImg,rowRange,colRange);

			//cv::Mat tempImg;
			//extendedImg.copyTo(tempImg);
			//cv::rectangle(tempImg,
			//	cv::Rect(colRangeLeft,rowRangeTop,colRangeRight-colRangeLeft,rowRangeBottom-rowRangeTop),
			//	cv::Scalar(255,255,255));
			//cv::imshow("sub",tempImg);
			//cv::waitKey(0);

			//cv::imshow("src",subSrcRegion);
			////cv::imshow("enhance",subEnhanceRegion);
			//cv::waitKey(0);
			this->interpolation(subSrcRegion,subEnhanceRegion,
							tileMappings[mapIndexTop*m_numTilesRow+mapIndexLeft],
							tileMappings[mapIndexBottom*m_numTilesRow+mapIndexLeft],
							tileMappings[mapIndexBottom*m_numTilesRow+mapIndexRight],
							tileMappings[mapIndexTop*m_numTilesRow+mapIndexRight]);

			//cv::imshow("src",extendedImg);
			//cv::imshow("enhance",extendedEnhanceImg);
			//cv::waitKey(0);
		}

	}
	//get the enhanceImg by removing the border of the extendedEnhanceImg
	cv::Mat tempImg(extendedEnhanceImg,
		            cv::Range(m_padRowPre,extendedRows-m_padRowPost),
					cv::Range(m_padColPre,extendedCols-m_padColPost));
	tempImg.copyTo(enhanceImg);

}

void MyCLAHE::extendImage(cv::Mat& srcImg,cv::Mat& extendedImg)
{
	CV_Assert(srcImg.channels()==1 && srcImg.depth()==CV_8U);

	unsigned int srcRows = srcImg.rows;	
	unsigned int srcCols = srcImg.cols;

	unsigned int padRow = 0;//the number of padded rows
	unsigned int padCol = 0;//the number of padded cols

	bool rowDiv = (srcRows % m_numTilesRow == 0);
	bool colDiv = (srcCols % m_numTilesCol == 0);

	if(!(rowDiv && colDiv))
	{
		if(!rowDiv)
		{
			m_rowsTile = floor(srcRows*1.0/m_numTilesRow) + 1;
			padRow = m_rowsTile*m_numTilesRow - srcRows;
		}
		else
			m_rowsTile = srcRows / m_numTilesRow;
		
		if (!colDiv)
		{
			m_colsTile = floor(srcCols*1.0/m_numTilesCol) + 1;
			padCol = m_colsTile*m_numTilesCol - srcCols;
		}
		else
			m_colsTile = srcCols / m_numTilesCol;
		
	}
	else
	{
		m_rowsTile = srcRows / m_numTilesRow;
		m_colsTile = srcCols / m_numTilesCol;
	}

	bool rowEven = (m_rowsTile % 2 == 0);
	bool colEven = (m_colsTile % 2 == 0);

	if (!(rowEven && colEven))
	{
		if (!rowEven)
		{
			padRow += m_numTilesRow;
			m_rowsTile += 1;
		}

		if (!colEven)
		{
			padCol += m_numTilesCol;
			m_colsTile += 1;
		}
	}

	m_padRowPre = floor(padRow/2.0);
	m_padRowPost = ceil(padRow/2.0);
	m_padColPre = floor(padCol/2.0);
	m_padColPost = ceil(padCol/2.0);

	unsigned int extendRows = srcRows + padRow;
	unsigned int extendCols = srcCols + padCol;
	extendedImg.release();
	extendedImg = cv::Mat::zeros(extendRows,extendCols,srcImg.type());
	//Copy the data of srcImg to extendedImg
	for (unsigned int i=0;i<srcRows;++i)
	{
		uchar* data1 = srcImg.ptr<uchar>(i);
		uchar* data2 = extendedImg.ptr<uchar>(i+m_padRowPre);
		for (unsigned int j=0;j<srcCols;++j)
		{
			data2[j+m_padColPre] = data1[j];
		}
	}
	//deal with the border pixels
	for (unsigned int i=0;i<m_padRowPre;++i)
	{
		extendedImg.row(2*m_padRowPre-1-i).copyTo(extendedImg.row(i));
	}
	unsigned int tempRow = m_padRowPre + srcRows;
	for (unsigned int i=0;i<m_padRowPost;++i)
	{
		extendedImg.row(tempRow-1-i).copyTo
			        (extendedImg.row(tempRow+i));
	}
	for (unsigned int i=0;i<m_padColPre;++i)
	{
		extendedImg.col(2*m_padColPre-1-i).copyTo(extendedImg.col(i));
	}
	unsigned int tempCol = m_padColPre + srcCols;
	for (unsigned int i=0;i<m_padColPost;++i)
	{
		extendedImg.col(tempCol-1-i).copyTo
			          (extendedImg.col(tempCol+i));
	}
}

void MyCLAHE::computeTileMappings(cv::Mat extendedImg,
	              std::vector<std::vector<unsigned long int>>& tileMappings)
{
	std::vector<unsigned long int> tileHist(numBins_CLAHE,0);

	//compute the cliplimit from the m_fCliplimit
	unsigned long int cliplimit = 0;
	if (m_fCilplimit>0.0 && m_fCilplimit<1.0)
	{
		unsigned long int minCliplimit = ceil(m_rowsTile*m_colsTile*1.0/numBins_CLAHE);
		cliplimit = minCliplimit + floor(m_fCilplimit*(m_rowsTile*m_colsTile-minCliplimit));
	} 
	else
	{
		cliplimit = 1UL<<14;//large number, do not clip
	}

	//Compute the tile mapping function
	for (unsigned int i=0;i<m_numTilesRow;++i)
	{
		cv::Range tileRowRange(i*m_rowsTile,(i+1)*m_rowsTile);
		for (unsigned int j=0;j<m_numTilesCol;++j)
		{
			cv::Range tileColRange(j*m_colsTile,(j+1)*m_colsTile);
			cv::Mat tileMat(extendedImg,tileRowRange,tileColRange);

			this->computeHist(tileMat,tileHist);

			//this->showHistogram(tileHist,"tileHist",cv::Scalar(255,0,0));
			//cv::waitKey(0);

			
			this->clipHistogram(tileHist,cliplimit);

			//this->showHistogram(tileHist,"clipHist",cv::Scalar(0,0,255));
			//cv::waitKey(0);

			std::vector<unsigned long int> tileMapping(numBins_CLAHE,0);
			this->mapHistogram(tileHist,tileMapping,m_rowsTile*m_colsTile);
			tileMappings.push_back(tileMapping);

			//this->showHistogram(tileMapping,"mapHist",cv::Scalar(0,255,0));
			//cv::waitKey(0);
		}
	}
}

void MyCLAHE::computeHist(cv::Mat& tileMat,
	                      std::vector<unsigned long int>& tileHist)
{
	unsigned int tileRows = tileMat.rows;
	unsigned int tileCols = tileMat.cols;

	tileHist.clear();
	for (unsigned int i=0;i<numBins_CLAHE;++i)
		tileHist.push_back(0L);

	for (unsigned int i=0;i<tileRows;++i)
	{
		uchar* data = tileMat.ptr<uchar>(i);
		for (unsigned int j=0;j<tileCols;++j)
		{
			uchar pixel = data[j];
			tileHist[pixel] += 1;
		}
	}
}

void MyCLAHE::clipHistogram(std::vector<unsigned long int>& tileHist,
	                        unsigned long int cliplimit)
{

	//this->showHistogram(tileHist,"hist",cv::Scalar(255,0,0));

	//Compute the number of excess pixels
	unsigned long int numExcessPixels = 0;
	long numBinExcess = 0;
	for (unsigned int i=0;i<numBins_CLAHE;++i)
	{
		numBinExcess = static_cast<long>(tileHist[i]) 
			           - static_cast<long>(cliplimit);
		if(numBinExcess>0)
			numExcessPixels += numBinExcess;
	}
	//Clip the tile histogram and redistribut the excess pixels in each bins
	unsigned long int numBinIncr = floor(numExcessPixels*1.0 / numBins_CLAHE);
	unsigned long int numUpper = cliplimit - numBinIncr;//Bins larger than numUpper set to cliplimit
	for (unsigned int i=0;i<numBins_CLAHE;++i)
	{
		if(tileHist[i] > cliplimit)
			tileHist[i] = cliplimit;
		else
		{
			if(tileHist[i] > numUpper)
			{
				tileHist[i] = cliplimit;
				//numExcessPixels -= (tileHist[i] - numUpper);
				numExcessPixels -= (cliplimit - tileHist[i]);
			}
			else
			{
				tileHist[i] += numBinIncr;
				numExcessPixels -= numBinIncr;
			}
		}
	}

	while (numExcessPixels > 0)
	{
		int startBin = 0;
		while (numExcessPixels > 0 && startBin < numBins_CLAHE)
		{
			int binStep = floor(numBins_CLAHE*1.0 / numExcessPixels);
			if(binStep<1 || binStep>=255) binStep = 1;
			for (int i=0;i<numBins_CLAHE && numExcessPixels!=0;i+=binStep)
			{
				if(tileHist[i] < cliplimit)
				{
					tileHist[i] += 1;
					numExcessPixels--;
				}
			}

			startBin++;
			if(startBin==numBins_CLAHE) startBin = 0;
		}
	}
	//this->showHistogram(tileHist,"hist",cv::Scalar(255,0,0));
	//cv::waitKey(0);
}

void MyCLAHE::mapHistogram(std::vector<unsigned long int>& tileHist,
	                       std::vector<unsigned long int>& mapHist,
						   unsigned long int numOfPixels)
{
	unsigned long int sum = 0;
	double _area = 255.0/numOfPixels;

	for (unsigned int i=0;i<numBins_CLAHE;++i)
	{
		sum += tileHist[i];
		mapHist[i] = sum*_area + 0.5;
		if(mapHist[i] > 255) mapHist[i] = 255;
	}
}

void MyCLAHE::interpolation(cv::Mat& srcImg,cv::Mat& enhanceImg,
	std::vector<unsigned long int>& mapHist_LeftTop,
	std::vector<unsigned long int>& mapHist_LeftBottom,
	std::vector<unsigned long int>& mapHist_RightBottom,
	std::vector<unsigned long int>& mapHist_RightTop)
{
	unsigned int Rows = srcImg.rows;
	unsigned int Cols = srcImg.cols;

	unsigned int RowCoef,RowInvCoef;
	unsigned int ColCoef,ColInvCoef;
	unsigned int NumPixels = Rows * Cols;
	double _NumPixels = 1.0/NumPixels;

	for (RowCoef=0,RowInvCoef=Rows;RowCoef<Rows;++RowCoef,--RowInvCoef)
	{
		uchar* data_src = srcImg.ptr<uchar>(RowCoef);
		uchar* data_enhance = enhanceImg.ptr<uchar>(RowCoef);
		for (ColCoef=0,ColInvCoef=Cols;ColCoef<Cols;++ColCoef,--ColInvCoef)
		{
			uchar pixel = data_src[ColCoef];
			//data_enhance[ColCoef] = 
			double data = 
				            (RowInvCoef*(ColInvCoef*mapHist_LeftTop[pixel]
			                 + ColCoef*mapHist_RightTop[pixel]) 
				             + RowCoef*(ColInvCoef*mapHist_LeftBottom[pixel]
							 + ColCoef*mapHist_RightBottom[pixel]))*_NumPixels;
			data_enhance[ColCoef] = static_cast<unsigned long int>(data);
		}
	}
}

void MyCLAHE::showHistogram(std::vector<unsigned long int> hist,
							std::string windowName,
							cv::Scalar color)
{
	assert(hist.size()!=0);

	//define the size of the histogram figure
	unsigned long int hist_height = 256;
	unsigned long int hist_width = hist.size();
	//Normalize the histogram
	unsigned long int minValue = hist[0];
	unsigned long int maxValue = hist[0];
	unsigned long int tempValue = 0;
	for (unsigned int i=0;i<hist.size();++i)
	{
		tempValue = hist[i];
		if(tempValue < minValue)
			minValue = tempValue;
		else
		{
			if(tempValue > maxValue)
				maxValue = tempValue;
		}
	}

	hist_height = maxValue;

	//double scale = 255.0 / (maxValue-minValue);
	//for (unsigned int i=0;i<hist.size();++i)
	//{
	//	tempValue = hist[i];
	//	tempValue = (tempValue-minValue)*scale + 0.5;
	//	hist[i] = tempValue;
	//}
	//draw the histogram figure
	cv::Mat histFigureMat = cv::Mat::zeros(hist_height,hist_width,CV_8UC3);
	for (unsigned int i=0;i<hist.size();++i)
	{
		cv::Point pt1(i,hist_height-1);
		cv::Point pt2(i,hist_height-1-hist[i]);
		cv::line(histFigureMat,pt1,pt2,color);
	}

	cv::imshow(windowName,histFigureMat);
	//cv::waitKey(0);
}


//template <class T>
//void showHistogram(std::vector<T> hist,
//	               std::string windowName,
//				   cv::Scalar color)
//{
//	assert(hist.size()!=0);
//
//	//define the size of the histogram figure
//	unsigned long int hist_height = 256;
//	unsigned long int hist_width = hist.size();
//	//Normalize the histogram
//	T minValue = hist[0];
//	T maxValue = hist[0];
//	T tempValue = 0;
//	for (unsigned int i=0;i<hist.size();++i)
//	{
//		tempValue = hist[i];
//		if(tempValue < minValue)
//			minValue = tempValue;
//		else
//		{
//			if(tempValue > maxValue)
//				maxValue = tempValue;
//		}
//	}
//	double scale = 255.0 / (maxValue-minValue);
//	for (unsigned int i=0;i<hist.size();++i)
//	{
//		tempValue = hist[i];
//		tempValue = (tempValue-minValue)*scale + 0.5;
//		hist[i] = tempValue;
//	}
//	//draw the histogram figure
//	cv::Mat histFigureMat = cv::Mat::zeros(hist_height,hist_width,CV_8UC3);
//	for (T i=0;i<hist.size();++i)
//	{
//		cv::Point pt1(i,hist_height-1);
//		cv::Point pt2(i,hist_height-1-hist[i]);
//		cv::line(histFigureMat,pt1,pt2,color);
//	}
//
//	cv::imshow(windowName,histFigureMat);
//	//cv::waitKey(0); 
//}